# run_bert.py
from transformers import BertTokenizer, BertForSequenceClassification
import torch
from colorama import Fore, init
import logging
import warnings

# 忽略 Flash Attention 警告
warnings.filterwarnings("ignore", category=UserWarning, message=".*Torch was not compiled with flash attention.*")

# 配置日志记录，不使用 encoding 参数
logger = logging.getLogger()
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(message)s', '%Y-%m-%d %H:%M:%S')

# 手动打开日志文件并指定编码
file_handler = logging.FileHandler('logs/detection.log', encoding='utf-8')
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)

init(autoreset=True)

# 加载模型和 tokenizer
model_path = "models/procedure_injection_bert"
model = BertForSequenceClassification.from_pretrained(model_path)
tokenizer = BertTokenizer.from_pretrained(model_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

label_map = {
    0: "✅ 正常 SQL",
    1: "🚨 SQL 注入 ⚠️",
    2: "🚨 存储过程注入 ⚠️"
}

# 存储过程关键字
STORED_PROCEDURE_KEYWORDS = [
    "EXEC",
    "EXECUTE",
    "sp_executesql",
    "DECLARE",
    "SET",
    "OPENROWSET",
    "BEGIN",
    "END",
    "CREATE PROCEDURE",
    "ALTER PROCEDURE",
    "DROP PROCEDURE",
    "CREATE FUNCTION",
    "ALTER FUNCTION",
    "DROP FUNCTION"
]

def clean_text(text):
    """去除存储过程关键字"""
    text_cleaned = text
    for keyword in STORED_PROCEDURE_KEYWORDS:
        text_cleaned = text_cleaned.replace(keyword, "")
        text_cleaned = text_cleaned.replace(keyword.lower(), "")
    return text_cleaned.strip()

print(Fore.CYAN + "\n🔍 英文 BERT SQL 注入实时检测器（含存储过程检测）已启动，输入 SQL 语句（输入 exit 退出）")

while True:
    try:
        query = input("SQL > ").strip()
        if query.lower() in ["exit", "quit", "q"]:
            print(Fore.YELLOW + "👋 已退出。")
            break
        if not query:
            continue

        # Step 1: 是否含存储过程关键字
        contains_proc_keyword = any(kw.lower() in query.lower() for kw in STORED_PROCEDURE_KEYWORDS)

        if contains_proc_keyword:
            # Step 2: 去除关键字
            cleaned_query = clean_text(query)

            # Step 3: 剩余部分为空 / 纯数字 直接认定正常
            if cleaned_query == "" or cleaned_query.replace(" ", "").isdigit():
                final_label = 0
            else:
                # 仍有内容，走模型判断
                inputs = tokenizer(cleaned_query, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
                with torch.no_grad():
                    outputs = model(**inputs)
                    pred_cleaned = torch.argmax(outputs.logits, dim=1).item()

                if pred_cleaned == 1:  # 剩余仍是注入
                    final_label = 2
                else:
                    final_label = 0  # 合法调用

        else:
            # Step 4: 普通流程直接判断
            inputs = tokenizer(query, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
            with torch.no_grad():
                outputs = model(**inputs)
                final_label = torch.argmax(outputs.logits, dim=1).item()

        # 输出结果
        result = label_map.get(final_label, "未知类别")
        color = Fore.GREEN if final_label == 0 else Fore.RED
        print(color + f"检测结果：{result}\n")

        # 记录日志
        logging.info(f"输入 SQL: {query}, 检测结果: {result}")

    except KeyboardInterrupt:
        print(Fore.YELLOW + "\n手动退出。")
        break